package org.example.brave.dubbo.provider.filter;

import brave.Span;
import brave.Tracer;
import brave.Tracing;
import brave.propagation.CurrentTraceContext;
import brave.propagation.Propagation;
import brave.propagation.TraceContext;
import brave.sampler.SamplerFunction;
import org.apache.dubbo.common.extension.Activate;
import org.apache.dubbo.rpc.*;
import zipkin2.reporter.AsyncReporter;
import zipkin2.reporter.brave.ZipkinSpanHandler;
import zipkin2.reporter.okhttp3.OkHttpSender;

import java.util.Map;

@Activate
public class BraveProviderFilter implements Filter {
    public static OkHttpSender sender;
    public static Tracing tracing;
    public static SamplerFunction<Invocation> sampler;

    static {
        sampler = new SamplerFunction<>() {
            @Override
            public Boolean trySample(Invocation invocation) {
                return !invocation.getMethodName().equalsIgnoreCase("getMetadataInfo");
            }
        };



        sender = OkHttpSender.newBuilder().endpoint("http://localhost:9411/api/v2/spans").build();
        tracing = Tracing.newBuilder().localServiceName("brave-dubbo-provider")
                .addSpanHandler(ZipkinSpanHandler.create(AsyncReporter.builder(sender).build()))
                .build();
    }

    Propagation.RemoteGetter<Map<String, Object>> GETTER = new Propagation.RemoteGetter<Map<String, Object>>() {
        @Override
        public Span.Kind spanKind() {
            return Span.Kind.SERVER;
        }

        @Override
        public String get(Map<String, Object> serverAttachment, String key) {
            return (String) serverAttachment.get(key);
        }
    };

    private boolean isConsumer(Invocation invocation) {
        Invoker<?> invoker = invocation.getInvoker();
        return invoker.getUrl().getParameter("side", "provider").equals("consumer");
    }

    @Override
    public Result invoke(Invoker<?> invoker, Invocation invocation) throws RpcException {
        System.out.println("Consumer judge whether provider or consumer: " + (isConsumer(invocation) ? "consumer" : "provider"));
        // recorde span... (learn as org.example.controller.ProviderController)
        Result result = null;
        CurrentTraceContext currentTraceContext = tracing.currentTraceContext();
        // ServerAttachment接收到客户端传递过来的参数
        Map<String, Object> serverAattachment = RpcContext.getServerAttachment().getObjectAttachments();
        TraceContext context = (TraceContext) serverAattachment.get(TraceContext.class.getName());
        if (context != null) {
            CurrentTraceContext.Scope scope = currentTraceContext.maybeScope(context);
            try {
                result = process(invoker, invocation);
            } catch (Exception e) {
                e.printStackTrace();
            } finally {
                scope.close();
            }
            return result;
        }

        TraceContext.Extractor<Map<String, Object>> extractor = tracing.propagation().extractor(GETTER);
        Tracer tracer = tracing.tracer();
        if (!sampler.trySample(invocation)) {
            return process(invoker, invocation);
        }
        Span span = tracer
                .nextSpan(extractor.extract(RpcContext.getServerAttachment().getObjectAttachments()))
                .name(invocation.getMethodName());
        // 1. Start the span and add trace headers to the request
        span.start();
        // 2. Put the span in scope so things like log integration works
        CurrentTraceContext.Scope scope = currentTraceContext.newScope(span.context());
        try  {
            result = process(invoker, invocation); // 3. Invoke the request
        } catch (Exception e) {
            span.error(e); // 4. Catch any errors
            e.printStackTrace();
        } finally {
            // 5. Complete the span
            span.kind(Span.Kind.SERVER);
            span.finish();
            scope.close();
        }
        return result;
    }

    private Result process(Invoker<?> invoker, Invocation invocation) {
        Result result = null;
        String clientIp = RpcContext.getContext().getRemoteHost();
        result = invoker.invoke(invocation);
        System.out.println("Remote IP: " + clientIp);
        return result;
    }
}
